package com.rizzo.analytics.services.web;

import org.apache.http.Header;
import org.apache.http.auth.Credentials;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.impl.auth.DigestScheme;
import org.apache.http.message.BasicHeader;
import org.apache.http.message.BasicHttpRequest;
import org.apache.http.protocol.BasicHttpContext;
import org.springframework.http.HttpMethod;
import org.springframework.http.HttpStatus;
import org.springframework.http.client.ClientHttpRequest;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.util.Assert;
import org.springframework.web.client.*;

import java.io.IOException;
import java.net.URI;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;

/**
 * DigestAuthRestTemplate is RestTemplate for Digest Authentication.
 * This class could be assigned to user session and re-used until user session is valid
 *
 * @author Reza Aliakbari(raliakbari@gmail.com)
 * @version 1.0.1, 5/20/12
 */
public class DigestAuthRestTemplate extends RestTemplate {

    protected static final String WWW_AUTHENTICATE = "WWW-Authenticate";

    // Wrap error handler for UnAuthorized error to get the challenge from response header
    ResponseErrorHandler errorHandler = new DigestErrorHandlerWrapper(new DefaultResponseErrorHandler());

    public void setErrorHandler(ResponseErrorHandler errorHandler) {
        Assert.notNull(errorHandler, "'errorHandler' must not be null");
        // Wrap error handler for UnAuthorized error to get the challenge from response header
        this.errorHandler = new DigestErrorHandlerWrapper(errorHandler);
    }

    /**
     * Return the error handler. By default, this is the {@link DefaultResponseErrorHandler}.
     */
    public ResponseErrorHandler getErrorHandler() {
        return this.errorHandler;
    }

    protected Credentials credentials;

    /**
     * hostChallengeCache caches challenges for hosts, each host may have multiple realms,
     * But this cache keeps the last challenge for the last call on the host
     */
    protected static ConcurrentHashMap<String, String> hostChallengeCache = new ConcurrentHashMap<String, String>();

    public DigestAuthRestTemplate(String username, String password) {
        this.credentials = new UsernamePasswordCredentials(username, password);
    }

    public DigestAuthRestTemplate(Credentials credentials) {
        this.credentials = credentials;
    }

    protected <T> T doExecute(URI url, HttpMethod method, RequestCallback requestCallback,
                              ResponseExtractor<T> responseExtractor) throws RestClientException {
        try {
            Header solution = createSolution(url, method);
            if (solution != null) {
                requestCallback = new DigestSolutionRequestCallbackWrapper(solution, requestCallback);
            }
            return super.doExecute(url, method, requestCallback, responseExtractor);
        } catch (DigestUnAuthorizedException digestUnAuthExp) {
            // create solution by the new challenge
            Header solution = createSolution(url, method, digestUnAuthExp.getChallenge());

            if (solution == null) {
                throw digestUnAuthExp;
            } else {
                DigestSolutionRequestCallbackWrapper dgRequestCallbackWrapper
                        = new DigestSolutionRequestCallbackWrapper(solution, requestCallback);
                // Call the API again with digest headers
                return super.doExecute(url, method, dgRequestCallbackWrapper, responseExtractor);
            }
        }
    }

    protected Header createSolution(URI uri, HttpMethod method) {
        return createSolution(uri, method, null);
    }

    protected Header createSolution(URI uri, HttpMethod method, String challenge) {

        if (challenge == null) {
            // get challenge from cache
            challenge = hostChallengeCache.get(uri.getHost());
            // If still challenge is null, return null
            if (challenge == null) return null;
        } else {
            // update challenge in cache
            hostChallengeCache.put(uri.getHost(), challenge);
        }

        Header challengeHeader = new BasicHeader(WWW_AUTHENTICATE, challenge);
        // TODO DigestScheme is not thread safe, we need a schema that compile WWW_AUTHENTICATE and keeps it as static variable and generate solution, The current solution could be a bit slow
        DigestScheme digestSolutionProvider = new DigestScheme();
        try {
            digestSolutionProvider.processChallenge(challengeHeader);
            return digestSolutionProvider
                    .authenticate(credentials, new BasicHttpRequest(method.name(), uri.getPath()), new BasicHttpContext());
        } catch (Exception e) {
            logger.error("couldn't parse new challenge " + challenge, e);
            return null;
        }
    }

    protected class DigestSolutionRequestCallbackWrapper implements RequestCallback {
        private Header solution;
        private RequestCallback wrappedRequestCallback;

        public DigestSolutionRequestCallbackWrapper(final Header solution, RequestCallback wrappedRequestCallback) {
            this.solution = solution;
            this.wrappedRequestCallback = wrappedRequestCallback;
        }

        public void doWithRequest(ClientHttpRequest request) throws IOException {
            // Add solution to header
            request.getHeaders().add(solution.getName(), solution.getValue());

            if (this.wrappedRequestCallback != null) {
                // Send request to the wrapped request call back
                this.wrappedRequestCallback.doWithRequest(request);
            }
        }
    }

    protected class DigestUnAuthorizedException extends HttpClientErrorException {

        private String challenge;

        public DigestUnAuthorizedException(String challenge) {
            super(HttpStatus.UNAUTHORIZED);
            this.challenge = challenge;
        }

        public String getChallenge() {
            return challenge;
        }
    }


    protected class DigestErrorHandlerWrapper implements ResponseErrorHandler {
        ResponseErrorHandler wrappedResponseErrorHandler;

        public DigestErrorHandlerWrapper(ResponseErrorHandler wrappedResponseErrorHandler) {
            Assert.notNull(wrappedResponseErrorHandler, "'wrappedResponseErrorHandler' must not be null");
            this.wrappedResponseErrorHandler = wrappedResponseErrorHandler;
        }

        public boolean hasError(ClientHttpResponse response) throws IOException {
            return response.getStatusCode().equals(HttpStatus.UNAUTHORIZED)
                    ||
                    this.wrappedResponseErrorHandler.hasError(response);
        }

        public void handleError(ClientHttpResponse response) throws IOException {

            HttpStatus statusCode = response.getStatusCode();
            if (statusCode.equals(HttpStatus.UNAUTHORIZED)) {
                List<String> wwwAuthHeaders = response.getHeaders().get(WWW_AUTHENTICATE);
                if (wwwAuthHeaders != null) {
                    String challenge = wwwAuthHeaders.get(0);
                    if (challenge != null) throw new DigestUnAuthorizedException(challenge);
                }
            }

            this.wrappedResponseErrorHandler.handleError(response);
        }
    }

}